from __future__ import absolute_import

from networks.AugmentCE2P import resnet101
from networks.parsenet50 import FaceParseNet101
from networks.AugmentCE2P_gcnet import resnet101_GCNet
from networks.AugmentCE2P_gcnet_with_gaussian import resnet101_GCNet_MASK
from networks.AugmentCE2P_gaussian import CE2P_gaussian
from networks.AugmentCE2P_BN import resnet101_BN
from networks.ehanet import EHANet18
from networks.AugmentCE2P_resnet_st import resnetst101
from networks.AugmentCE2P_gcnet_BN import resnet101_GCNet_BN
from networks.AugmentCE2P_OCR import resnet101_ocr
from networks.AugmentCE2P_OCR_2 import resnet101_ocr2
from networks.AugmentCE2P_video import CE2P_video
from networks.seg_hrnet_ocr import hrnet_ocr
__factory = {
    'resnet101': resnet101,
    'FaceParseNet101':FaceParseNet101,
    'resnet101_GCNet':resnet101_GCNet,
    'resnet101_GCNet_MASK':resnet101_GCNet_MASK,
    'resnet101_BN':resnet101_BN,
    "resnetst101":resnetst101,
    'EHANet18':EHANet18,
    'CE2P_gaussian':CE2P_gaussian,
    'resnet101_GCNet_BN':resnet101_GCNet_BN,
    'resnet101_ocr':resnet101_ocr,
    'CE2P_video':CE2P_video,
    'resnet101_ocr2':resnet101_ocr2,
    'hrnet_ocr':hrnet_ocr
}
def init_model(name, *args, **kwargs):
    if name not in __factory.keys():
        raise KeyError("Unknown model arch: {}".format(name))
    return __factory[name](*args, **kwargs)